#include "Drivetrain.h"

void Drivetrain::begin(Adafruit_MotorShield &shield) {
  _flMotor = shield.getMotor(1);
  _frMotor = shield.getMotor(2);
  _blMotor = shield.getMotor(3);
  _brMotor = shield.getMotor(4);
  stopAllMotors();
}

void Drivetrain::applyForward(int speed) {
  speed = constrain(speed, 0, 255);
  _flMotor->setSpeed(speed); _flMotor->run(FORWARD);
  _frMotor->setSpeed(speed); _frMotor->run(FORWARD);
  _blMotor->setSpeed(speed); _blMotor->run(FORWARD);
  _brMotor->setSpeed(speed); _brMotor->run(FORWARD);
}

void Drivetrain::applyBackward(int speed) {
  speed = constrain(speed, 0, 255);
  _flMotor->setSpeed(speed); _flMotor->run(BACKWARD);
  _frMotor->setSpeed(speed); _frMotor->run(BACKWARD);
  _blMotor->setSpeed(speed); _blMotor->run(BACKWARD);
  _brMotor->setSpeed(speed); _brMotor->run(BACKWARD);
}

void Drivetrain::applyTurn(int speed) {
  speed = constrain(speed, 0, 255);
  // Pivot turn: left/right sides driven in opposite directions.
  uint8_t leftDir  = (_heading > 0) ? BACKWARD : FORWARD;
  uint8_t rightDir = (_heading > 0) ? FORWARD  : BACKWARD;
  _flMotor->setSpeed(speed); _flMotor->run(leftDir);
  _blMotor->setSpeed(speed); _blMotor->run(leftDir);
  _frMotor->setSpeed(speed); _frMotor->run(rightDir);
  _brMotor->setSpeed(speed); _brMotor->run(rightDir);
}

void Drivetrain::stopAllMotors() {
  _flMotor->setSpeed(0); _flMotor->run(RELEASE);
  _frMotor->setSpeed(0); _frMotor->run(RELEASE);
  _blMotor->setSpeed(0); _blMotor->run(RELEASE);
  _brMotor->setSpeed(0); _brMotor->run(RELEASE);
}

void Drivetrain::startForwardHop(unsigned long runDurationMs, float speedScale) {
  if (isBusy()) return; // ignore requests while a maneuver is already in progress
  _targetSpeed = (int)(Motion::MOTOR_SPEED_BASE * speedScale);
  _plannedRunMs = runDurationMs;
  _state = ManeuverState::RAMP_UP;
  _stateStartMs = millis();
}

void Drivetrain::startAvoidanceManeuver(float speedScale) {
  if (isBusy()) return;
  _targetSpeed = (int)(Motion::MOTOR_SPEED_BASE * speedScale);
  _turnDurationMs = random(Motion::TURN_MS_MIN, Motion::TURN_MS_MAX + 1);
  _state = ManeuverState::BACKING_UP;
  _stateStartMs = millis();
}

void Drivetrain::emergencyStop() {
  stopAllMotors();
  _state = ManeuverState::IDLE;
}

void Drivetrain::update(unsigned long nowMs) {
  unsigned long elapsed = nowMs - _stateStartMs;

  switch (_state) {
    case ManeuverState::IDLE:
      break;

    case ManeuverState::RAMP_UP: {
      float frac = (Motion::RAMP_MS == 0) ? 1.0f : min(1.0f, (float)elapsed / (float)Motion::RAMP_MS);
      applyForward((int)(frac * _targetSpeed));
      if (frac >= 1.0f) { _state = ManeuverState::RUNNING; _stateStartMs = nowMs; }
      break;
    }

    case ManeuverState::RUNNING:
      if (elapsed >= _plannedRunMs) { _state = ManeuverState::RAMP_DOWN; _stateStartMs = nowMs; }
      break;

    case ManeuverState::RAMP_DOWN: {
      float frac = (Motion::RAMP_MS == 0) ? 1.0f : min(1.0f, (float)elapsed / (float)Motion::RAMP_MS);
      applyForward((int)((1.0f - frac) * _targetSpeed));
      if (frac >= 1.0f) { stopAllMotors(); _state = ManeuverState::IDLE; }
      break;
    }

    case ManeuverState::BACKING_UP:
      applyBackward(_targetSpeed);
      if (elapsed >= Motion::BACKUP_MS) { _state = ManeuverState::TURNING; _stateStartMs = nowMs; }
      break;

    case ManeuverState::TURNING:
      applyTurn(_targetSpeed);
      if (elapsed >= _turnDurationMs) {
        stopAllMotors();
        _heading = -_heading;
        _state = ManeuverState::IDLE;
      }
      break;
  }
}
